import time
import pandas as pd
import os
from openpyxl import Workbook


def init_xlsx(hive_file_name, mysql_file_name):
    """初始化csv_compare.xlsx文件"""
    wb = Workbook()
    ws = wb.active
    ws.cell(3, 1, f"{hive_file_name}")
    ws.cell(4, 1, f"{mysql_file_name}")
    return wb, ws, time.strftime('%y%m%d-%H%M%S', time.localtime(time.time()))


def sort_file(csv_reader, sorted_file_path):
    """
    对CSV所有列排序，并将结果保存到新的csv文件
    :param csv_reader:
    :param sorted_file_path: 排序后到CSV文件路径
    :return:
    """
    csv_reader.sort_values(list(csv_reader.columns.values), inplace=True, ignore_index=True)
    print(f'排序后的文件路径：{sorted_file_path}')
    csv_reader.to_csv(sorted_file_path, mode='w+', index=False)


def sort_file_new(file_path):
    csv_reader = pd.read_csv(file_path, low_memory=False)
    csv_reader.sort_values(list(csv_reader.columns.values), inplace=True, ignore_index=True)
    sorted_file_path = '_sort'.join(os.path.splitext(file_path))
    print(f'排序后的文件路径：{sorted_file_path}')
    csv_reader.to_csv(sorted_file_path, mode='w+', index=False)
    return sorted_file_path


def csv_compare(hive_file_path, mysql_file_path, rows=0, sort=False):
    """
    对比2个csv文件是否一致
    :param hive_file_path:
    :param mysql_file_path:
    :param rows: 若仅想对比前n行，则rows=n，默认rows=0为对比所有行
    :param sort: 是否需要排序后再做对比
    :return:
    """
    base_path, hive_file_name = os.path.split(hive_file_path)
    base_path, mysql_file_name = os.path.split(mysql_file_path)

    hive_reader = pd.read_csv(hive_file_path, low_memory=False)
    mysql_reader = pd.read_csv(mysql_file_path, low_memory=False)

    if sort:
        hive_file_name = '_sort'.join(os.path.splitext(hive_file_path))
        mysql_file_name = '_sort'.join(os.path.splitext(mysql_file_path))
        sort_file(hive_reader, hive_file_name)
        sort_file(mysql_reader, mysql_file_name)

    print(f"{hive_file_name} 的内容大小为: {hive_reader.shape}")
    print(f"{mysql_file_name} 的内容大小为: {mysql_reader.shape}")

    wb, ws, time_info = init_xlsx(hive_file_name, mysql_file_name)

    cur_t = time.time()

    temp_col = 2

    result_flag = True

    for j, col_name in enumerate(hive_reader.columns):
        print(f"第 {j + 1} 列数据比较...")
        if rows:
            diff = hive_reader[col_name].head(rows).compare(mysql_reader[mysql_reader.columns[j]].head(rows))
        else:
            diff = hive_reader[col_name].compare(mysql_reader[mysql_reader.columns[j]])

        if diff.shape[0]:
            ws.cell(1, temp_col, col_name)

            result_flag = False
            temp_row = 2

            for index, row in diff.iterrows():
                ws.cell(temp_row, temp_col, f"第 {index + 1} 行对比结果")
                ws.cell(temp_row + 1, temp_col, row[0])
                ws.cell(temp_row + 2, temp_col, row[1])
                temp_row += 4
            temp_col += 1

    print(f"对比结果为：{'完全一致' if result_flag else '不一致'}")
    print(f'耗时: {time.time() - cur_t}')
    print(f'将对比结果存储至csv_compare_{time_info}.xlsx ...')
    wb.save(f"/Users/xunjie.yin/PycharmProjects/pythonProject/csv_compare/csv_compare_{time_info}.xlsx")


if __name__ == '__main__':
    # hive_file_path = "/Users/xunjie.yin/Downloads/query-hive-110773.csv"
    # mysql_file_path = "/Users/xunjie.yin/Downloads/rep_fin_reg_db_master_tabungan_ss_d_2022-08-01T13-12-00.csv"
    # csv_compare(hive_file_path, mysql_file_path, rows=0, sort=True)
    sort_file_new('/Users/xunjie.yin/Downloads/query-hive-110773.csv')

